#include "memorydevice.h"
#include "misc.h"
#include "logger.h"
#include <vector>

#include "error.h"

CVardata::CVardata()
{
    iData = 0;
    fData = 0.0;
    cData = 0;
	sData = 0;
	szData = "";
    type = T_INT;

    varcount = 1;
}


std::string MemoryDevice::readString(HANDLE process, unsigned long address, int &err,
    unsigned int len)
{
    std::string fullstr;
    unsigned char buffer = 0;
    SIZE_T bytesread;
    err = 0;
    int success;

    long stroffset = 0;
    unsigned int count = 0;
    while( 1 ) // read until we hit a NULL
    {
        success = ReadProcessMemory(process, (LPCVOID)(address + stroffset),
            (void*)&buffer, sizeof(unsigned char), &bytesread);
        stroffset++;

        fullstr.push_back(buffer);
        count++;

        if( !len && buffer == '\0' )
            break;

        if( len && count >= len )
        {
            fullstr.push_back('\0');
            break;
        }
    }

    if( success == 0 )
        err = MEMORY_READ_FAIL;

    return fullstr;
}

std::wstring MemoryDevice::readUString(HANDLE process, unsigned long address, int &err,
    unsigned int len)
{
    std::wstring fullstr;
    wchar_t buffer = 0;
    SIZE_T bytesread;
    err = 0;
    int success;


    long stroffset = 0;
    unsigned int count = 0;
    while( 1 ) // read until we hit a NULL
    {
        success = ReadProcessMemory(process, (LPCVOID)(address + stroffset),
            (void*)&buffer, sizeof(wchar_t), &bytesread);
        stroffset += sizeof(wchar_t);

        fullstr.push_back(buffer);
        count++;

        if( buffer == 0 )
            break;

        if( len != 0 && count >= len )
        {
            fullstr.push_back('\0');
            break;
        }
    }

    if( success == 0 )
        err = MEMORY_READ_FAIL;

    return fullstr;
}


std::string MemoryDevice::readStringPtr(HANDLE process, unsigned long address,
    long offset, int &err, unsigned int len)
{
    std::string fullstr;
    char buffer = 0;
    SIZE_T bytesread;
    err = 0;
    int success;
    long real_addr = 0;

    // read the value of the pointer, first
    success = ReadProcessMemory(process, (LPCVOID)address,
        (void*)&real_addr, sizeof(long), &bytesread);

    if( success == 0 ) {
        err = MEMORY_READ_FAIL;
        return 0; }

    long stroffset = 0;
    unsigned int count = 0;
    while( 1 )
    {
        success = ReadProcessMemory(process,
            (LPCVOID)(real_addr + offset + stroffset), (void*)&buffer,
            sizeof(char), &bytesread);
        stroffset += sizeof(char);

        fullstr.push_back(buffer);
        count++;

        if( !len && buffer == '\0' )
            break;

        if( len && count >= len )
        {
            fullstr.push_back('\0');
            break;
        }
    }

    if( success == 0 )
        err = MEMORY_READ_FAIL;

    return fullstr;
}

std::wstring MemoryDevice::readUStringPtr(HANDLE process, unsigned long address,
    long offset, int &err, unsigned int len)
{
    std::wstring fullstr;
    wchar_t buffer = 0;
    SIZE_T bytesread;
    err = 0;
    int success;
    long real_addr = 0;

    // read the value of the pointer, first
    success = ReadProcessMemory(process, (LPCVOID)address,
        (void*)&real_addr, sizeof(long), &bytesread);

    if( success == 0 ) {
        err = MEMORY_READ_FAIL;
        return 0; }

    long stroffset = 0;
    unsigned int count = 0;
    while( 1 )
    {
        success = ReadProcessMemory(process,
            (LPCVOID)(real_addr + offset + stroffset), (void*)&buffer,
            sizeof(wchar_t), &bytesread);
        stroffset += sizeof(wchar_t);

        fullstr.push_back(buffer);
        count++;

        if( !len && buffer ==0 )
            break;

        if( len && count >= len )
        {
            fullstr.push_back('\0');
            break;
        }
    }

    if( success == 0 )
        err = MEMORY_READ_FAIL;

    return fullstr;
}

void MemoryDevice::writeString(HANDLE process, unsigned long address, char *data, int &err, unsigned int len)
{
	SIZE_T byteswritten = 0;
	err = 0;
	int success = 0;
	DWORD old;

	VirtualProtectEx(process, (void *)address, (size_t)len, PAGE_READWRITE, &old);
	success = WriteProcessMemory(process, (void *)address,
	(void*)data, (size_t)len, &byteswritten);
	VirtualProtectEx(process, (void *)address, (size_t)len, old, &old);

	if( success == 0 )
		err = MEMORY_WRITE_FAIL;
}

void MemoryDevice::writeStringPtr(HANDLE process, unsigned long address, long offset, char *data,
	int &err, unsigned int len)
{
	long real_addr = 0;
	SIZE_T bytesread = 0;
	SIZE_T byteswritten = 0;
	int success = 0;
	err = 0;
	DWORD old;

	// find the real (base) address to write to
	success = ReadProcessMemory(process, (LPCVOID)address,
	(void*)&real_addr, sizeof(long), &bytesread);

	if( success == 0 ) {
		err = MEMORY_READ_FAIL;
		return; }

	// now write to that address + offset
	VirtualProtectEx(process, (void *)(real_addr + offset), (size_t)len, PAGE_READWRITE, &old);
	success = WriteProcessMemory(process,
		(void *)(real_addr + offset), (void*)data,
		(size_t)len, &byteswritten);
	VirtualProtectEx(process, (void *)(real_addr + offset), (size_t)len, old, &old);

	if( success == 0 )
		err = MEMORY_WRITE_FAIL;
}


unsigned int MemoryDevice::readBatch_parsefmt(char *fmt, std::vector<CVardata> &out)
{
	unsigned int length = 0;
	out.clear();

	CVardata job;
	char c;

	for(unsigned int i = 0; i < strlen(fmt); i++)
	{
		c = fmt[i];

		// See if we are setting the varcount
		if( c >= '0' && c <= '9' )
		{
			// Read the number as a string, convert to int, advance
			char buffer[16];
			memset(&buffer, 0, 16);
			for(unsigned int j = 0; j < 15; j++)
			{
				char p = fmt[i + j];
				if( p < '0' || p > '9' )
				{ // End
					buffer[j] = 0;
					job.varcount = atoi((char*)&buffer);
					i += j - 1;
					break;
				}

				buffer[j] = p;
			}

			continue;
		}
		else // Setting a variable type
		{
			switch(c)
			{
				case 'b':
					job.type = T_BYTE;
					length += sizeof(char) * job.varcount;
					out.push_back(job);
				break;
				case 'B':
					job.type = T_UBYTE;
					length += sizeof(unsigned char) * job.varcount;
					out.push_back(job);
				break;
				case 's':
					job.type = T_SHORT;
					length += sizeof(short) * job.varcount;
					out.push_back(job);
				break;
				case 'S':
					job.type = T_USHORT;
					length += sizeof(unsigned short) * job.varcount;
					out.push_back(job);
				break;
				case 'i':
					job.type = T_INT;
					length += sizeof(int) * job.varcount;
					out.push_back(job);
				break;
				case 'I':
					job.type = T_UINT;
					length += sizeof(unsigned int) * job.varcount;
					out.push_back(job);
				break;
				case 'f':
					job.type = T_FLOAT;
					length += sizeof(float) * job.varcount;
					out.push_back(job);
				break;
				case 'c':
					job.type = T_SZSTRING;
					length += sizeof(char) * job.varcount;
					out.push_back(job);
				break;
				case '_':
					job.type = T_UNUSED;
					length += sizeof(char) * job.varcount;
					out.push_back(job);
				break;
			}

			job.varcount = 1; // Reset count
		}
	}

	return length;
}

/* Reads a block of memory, using 'fmt' as the format.
	Any read variables are returned via 'out'.
*/
void MemoryDevice::readBatch(HANDLE process, unsigned long address,
	char *fmt, std::vector<CVardata> &out, int &err)
{
	unsigned int readLen = 0;
	unsigned int cursorPos = 0;
	std::vector<CVardata> jobs;
	err = 0;
	out.clear();

	/* Parse format, get length */
	readLen = readBatch_parsefmt(fmt, jobs);

	/* Read memory */
	char *readBuffer = new char[readLen+1];
	if( readBuffer == NULL )
		allocationError();

	SIZE_T bytesread = 0;
	int success = 0;

	if( readBuffer == NULL )
		allocationError();

	success = ReadProcessMemory(process, (LPCVOID)address,
		(void *)readBuffer, readLen, &bytesread);

	if( success == 0 || bytesread != readLen )
	{
		err = MEMORY_READ_FAIL;
		delete []readBuffer;
		return;
	}

	/* Extract data, dump into output vector */
	int *iRead;
	float *fRead;
	char *cRead;
	short *sRead;
	CVardata outjob;

	for(unsigned int i = 0; i < jobs.size(); i++)
	{
		switch(jobs.at(i).type)
		{
			case T_BYTE:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					cRead = (char *)&readBuffer[cursorPos];
					outjob.cData = *cRead;
					outjob.type = T_BYTE;
					cursorPos += sizeof(char);
					out.push_back(outjob);
				}
			break;
			case T_UBYTE:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					cRead = (char *)&readBuffer[cursorPos];
					outjob.cData = *cRead;
					outjob.type = T_BYTE;
					cursorPos += sizeof(unsigned char);
					out.push_back(outjob);
				}
			break;
			case T_SHORT:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					sRead = (short *)&readBuffer[cursorPos];
					outjob.sData = *sRead;
					outjob.type = T_SHORT;
					cursorPos += sizeof(short);
					out.push_back(outjob);
				}
			break;
			case T_USHORT:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					sRead = (short *)&readBuffer[cursorPos];
					outjob.sData = *sRead;
					outjob.type = T_USHORT;
					cursorPos += sizeof(unsigned short);
					out.push_back(outjob);
				}
			break;
			case T_INT:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					iRead = (int *)&readBuffer[cursorPos];
					outjob.iData = *iRead;
					outjob.type = T_INT;
					cursorPos += sizeof(int);
					out.push_back(outjob);
				}
			break;
			case T_UINT:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					iRead = (int *)&readBuffer[cursorPos];
					outjob.iData = *iRead;
					outjob.type = T_UINT;
					cursorPos += sizeof(unsigned int);
					out.push_back(outjob);
				}
			break;
			case T_FLOAT:
				for(unsigned int j = 0; j < jobs.at(i).varcount; j++)
				{
					fRead = (float *)&readBuffer[cursorPos];
					outjob.fData = *fRead;
					outjob.type = T_FLOAT;
					cursorPos += sizeof(float);
					out.push_back(outjob);
				}
			break;
			case T_SZSTRING:
				{
					int len = jobs.at(i).varcount;
					// extra to ensure NULL terminator
					char *cBuff = new char[len + 2];

					if( cBuff == NULL )
						allocationError();

					memset(cBuff, 0, len + 1);
					strncpy(cBuff, (char *)&readBuffer[cursorPos], len);
					cursorPos += len;
					outjob.szData = cBuff;
					outjob.type = T_SZSTRING;
					out.push_back(outjob);

					delete []cBuff;
				}
			break;
			case T_UNUSED:
				cursorPos += sizeof(char) * jobs[i].varcount;
			break;
		}
	}

	delete []readBuffer;
}
